#include "limbo/arp/arp_request.hpp"
#include "limbo/arp/arp_response.hpp"
#include "limbo/arp/state.hpp"
#include "limbo/ethernet/state.hpp"
#include "limbo/stack.hpp"
#include "test-utils.h"

using namespace limbo;

using namespace limbo;

using Eth = ethernet::State<void>;
using Arp = arp::State<Eth>;
using MyStack = Stack<Layer<Arp>, Layer<Eth>>;
using Mac = ethernet::MacAddress;

uint8_t example_raw[]{
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x48, 0xba, 0x4e, 0x51, 0x39, 0xbb,
    0x08, 0x06, 0x00, 0x01, 0x08, 0x00, 0x06, 0x04, 0x00, 0x01, 0x48, 0xba,
    0x4e, 0x51, 0x39, 0xbb, 0xc0, 0xa8, 0x65, 0x01, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0xc0, 0xa8, 0x65, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};
auto example_chunk = Chunk(example_raw, sizeof(example_raw));

uint8_t src_raw[] = {0x48, 0xba, 0x4e, 0x51, 0x39, 0xbb};
uint8_t dst_raw[] = {0x48, 0xba, 0x4e, 0x51, 0x39, 0xbc};

auto ip_src = make_address("192.168.101.1");
auto ip_dst = make_address("192.168.101.2");
auto mac_src = ethernet::MacAddress(src_raw);
auto mac_dst = ethernet::MacAddress(dst_raw);

TEST_CASE("send", "[ethernet+arp]") {
  auto stack = MyStack();

  auto eth_ctx =
      Eth::Context{mac_src, Mac::broadcast, ethernet::EtherType::arp, nullptr};
  auto arp_ctx = Arp::Context{
      ip_src, ip_dst, mac_src, Mac::unknown, arp::Operation::request, &eth_ctx};
  unsigned char buff[sizeof(example_raw)];
  auto buff_chunk = Chunk(buff, sizeof(buff));

  SECTION("send all") {
    auto result = stack.send(buff_chunk, arp_ctx, {});
    REQUIRE(result);

    auto res_buff = result.consumed();
    CHECK((void *)res_buff.data() == (void *)buff);
    CHECK(res_buff.size() == example_chunk.size());
    CHECK(res_buff == example_chunk);
  }
  SECTION("send envelope") {
    auto sender = stack.send_envelope(buff_chunk, arp_ctx, {});
    REQUIRE(sender);
    CHECK(sender.header() == Chunk(example_raw, sender.header().size()));
    CHECK(sender.trailer() == Chunk());
  }
}

TEST_CASE("receive", "[ethernet+arp]") {
  auto stack = MyStack();
  auto &arp_state = stack.get<0, 0>();
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &arp_state);
  auto &p = arp_state.get_parsed();
  CHECK(p.operation == (uint16_t)arp::Operation::request);
  CHECK(p.sender_mac == mac_src);
  CHECK(p.target_mac == Mac::unknown);
  CHECK(p.sender_ip == ip_src);
  CHECK(p.target_ip == ip_dst);

  REQUIRE(p.container);
  CHECK(p.container->source == mac_src);
  CHECK(p.container->destination == Mac::broadcast);
  CHECK(p.container->type == (uint16_t)ethernet::EtherType::arp);
}

TEST_CASE("request/response", "[ethernet+arp]") {
  using Request = arp::ArpRequest<Arp>;
  using Response = arp::ArpResponse<Arp>;
  using MyStack = Stack<Layer<Request, Response>, Layer<Arp>, Layer<Eth>>;

  unsigned char buff[sizeof(example_raw)];
  auto buff_chunk = Chunk(buff, sizeof(buff));

  auto stack = MyStack();

  auto &request = stack.get<0, 0>();
  auto &response = stack.get<0, 1>();
  request.init(mac_src, ip_src);

  auto result = request.send(stack, buff_chunk, ip_dst);
  REQUIRE(result);

  auto req_buff = result.consumed();
  CHECK((void *)req_buff.data() == (void *)buff);
  CHECK(req_buff.size() == example_chunk.size());
  CHECK(req_buff == example_chunk);

  request.init(mac_src, ip_dst);
  result = stack.recv(buff_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == buff_chunk);
  REQUIRE(result.state() == &request);

  auto &req_packet = request.get_parsed();
  CHECK(req_packet.operation == (uint16_t)arp::Operation::request);
  CHECK(req_packet.sender_mac == mac_src);
  CHECK(req_packet.target_mac == Mac::unknown);
  CHECK(req_packet.sender_ip == ip_src);
  CHECK(req_packet.target_ip == ip_dst);

  response.init(mac_dst, ip_dst);
  result = response.send(stack, buff_chunk, request);
  REQUIRE(result);

  auto res_buff = result.consumed();
  CHECK((void *)res_buff.data() == (void *)buff);
  CHECK(res_buff.size() == example_chunk.size());

  response.init(mac_src, ip_src);
  result = stack.recv(res_buff, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == buff_chunk);
  REQUIRE(result.state() == &response);

  auto &res_packet = response.get_parsed();
  CHECK(res_packet.operation == (uint16_t)arp::Operation::response);
  CHECK(res_packet.sender_mac == mac_dst);
  CHECK(res_packet.target_mac == mac_src);
  CHECK(res_packet.sender_ip == ip_dst);
  CHECK(res_packet.target_ip == ip_src);
}
